!
! Copyright (C) 2000-2013 D. Sangalli, C. Attaccalite and the YAMBO team 
!              https://code.google.com/p/rocinante.org
! 
! This file is distributed under the terms of the GNU 
! General Public License. You can redistribute it and/or 
! modify it under the terms of the GNU General Public 
! License as published by the Free Software Foundation; 
! either version 2, or (at your option) any later version.
!
! This program is distributed in the hope that it will 
! be useful, but WITHOUT ANY WARRANTY; without even the 
! implied warranty of MERCHANTABILITY or FITNESS FOR A 
! PARTICULAR PURPOSE.  See the GNU General Public License 
! for more details.
!
! You should have received a copy of the GNU General Public 
! License along with this program; if not, write to the Free 
! Software Foundation, Inc., 59 Temple Place - Suite 330,Boston, 
! MA 02111-1307, USA or visit http://www.gnu.org/copyleft/gpl.txt.
!
subroutine fix_WFs_and_E(E,k,k_save,kpoints_map,old_nsym,S_contains_TR,action_kind)
 !
 use pars,                ONLY:SP,lchlen
 use com,                 ONLY:msg,warning,core_io_path,more_io_path
 use drivers,             ONLY:l_sc_run
 use memory_m,            ONLY:mem_est
 use IO_m,                ONLY:io_control,OP_WR_CL,NONE,OP_APP_WR_CL,OP_WR,OP_RD,VERIFY,&
&                              OP_APP,cp_file,rm_file,cp_directory,Fragmented_IO, &
&                              io_fragmented,OP_RD_CL,RD_CL,RD,DUMP,RD_CL_IF_END,WR_CL_IF_END
 use wave_func,           ONLY:ioWF,wf_nc_k, wf_igk,wf_ncx,wf_nb_io,wf_ng, &
 &                             wf_state,WF_load,wf_state,wf,WF_free,wf_nb_io_groups
 use electrons,           ONLY:levels,E_reset,n_spin,n_sp_pol
 use timing,              ONLY:live_timing
 use stderr,              ONLY:intc
 use R_lattice,           ONLY:bz_samp,g_rot,ng_closed,nkibz,ng_vec
 use D_lattice,           ONLY:spin_sop
 !
 implicit none
 !
 type(levels),intent(in)  :: E
 !
 type(bz_samp),intent(in) :: k
 type(bz_samp),intent(in) :: k_save
 integer,intent(in)       :: kpoints_map(2,k%nibz)
 !
 integer,intent(in)       :: old_nsym
 logical,intent(in)       :: S_contains_TR(old_nsym)
 !
 integer,intent(in)       :: action_kind
 !
 ! Work space
 !
 type(levels)          :: E_new
 !
 character(lchlen)     :: core_io_path_save,fragment_name
 integer               :: n_steps,nb1,nb2,ng_closed_save
 integer               :: ACTION_,ID
 integer               :: ierr,io_err
 !
 integer               :: ng_vec_save,ng_vec_new
 integer               :: wf_ng_save,wf_ng_new
 integer               :: wf_ncx_save,wf_ncx_new
 integer               :: wf_nc_k_save(k_save%nibz),wf_nc_k_new(k%nibz)
 integer,allocatable   :: wf_igk_save(:,:),wf_igk_new(:,:)
 !
 integer,     allocatable :: ic_rot_table(:) 
 real(SP),    allocatable :: wf_disk(:,:,:,:)
 complex(SP), allocatable :: wf_tmp(:,:)
 !
 ! Dummies
 !
 integer               :: is,ik,ik_save
 integer               :: ib,ib_grp,i_wf,i_spin
 integer               :: ic,ic_rot,ig,ig_rot
 !
 ! External functions
 !
 integer, external :: ioDB1
 !
 !
 call io_control(ACTION=OP_RD_CL,COM=NONE,SEC=(/1/),ID=ID)
 io_err=ioWF(ID)
 Fragmented_IO=io_fragmented(ID)
 !
 core_io_path_save=core_io_path
 core_io_path=more_io_path
 !
 call section('=',"Energies and WFs shells")
 !===========================
 !
 allocate(wf_igk_save(wf_ncx,k_save%nibz))
 wf_nc_k_save=wf_nc_k
 wf_igk_save =wf_igk
 wf_ncx_save =wf_ncx
 wf_ng_save  =wf_ng
 ng_vec_save =ng_vec
 !
 ! Rotate spheres of G-vectors centered on kpts
 call fix_WFs_Gshells(k,k_save,kpoints_map,old_nsym,action_kind)
 !
 allocate(wf_igk_new(wf_ncx,k%nibz))
 wf_nc_k_new=wf_nc_k
 wf_igk_new =wf_igk
 wf_ncx_new =wf_ncx
 wf_ng_new  =wf_ng
 ng_vec_new =ng_vec
 !
 if(wf_ncx_save/=wf_ncx_new) call msg('s',':: wf_ncx was reduced to avoid components in open shells')
 if(ng_vec_save/=ng_vec_new) call msg('s',':: ng_vec was increased to close the G-shell')
 !
 call map_energies(E,E_new,k,kpoints_map)
 !
 ! For some reason ioDB1 changes the value of ng_closed
 ! This is a work-around to this problem
 ng_closed_save=ng_closed
 !
 call msg('s',':: I/O...')
 call io_control(ACTION=OP_WR_CL,COM=NONE,SEC=(/1,2/),ID=ID)
 io_err=ioDB1(E_new,k,ID)
 call msg('l','done')
 !
 ng_closed=ng_closed_save
 ! 
 call section('=',"Wavefunctions")
 !===============================
 !
 ! After I/O of DB1 I need to restore the old WFs shells
 ! to make sure that WF_load works properly
 deallocate(wf_nc_k,wf_igk)
 allocate(wf_nc_k(k_save%nibz),wf_igk(wf_ncx_save,k_save%nibz))
 !
 wf_nc_k=wf_nc_k_save
 wf_igk =wf_igk_save
 wf_ncx =wf_ncx_save
 wf_ng  =wf_ng_save 
 !
 select case(action_kind)
 !
 case(1)
   !
   if (wf_ncx_save==wf_ncx_new) then
     !
     call msg('s',':: Copying existing database ...')
     !
     call cp_file(trim(core_io_path_save)//"/SAVE/s.wf",trim(more_io_path)//"/SAVE",ierr)
     call cp_file(trim(core_io_path_save)//"/SAVE/ns.wf",trim(more_io_path)//"/SAVE",ierr)
     do ik=1,k_save%nibz
       do ib_grp=1,wf_nb_io_groups
         fragment_name='s.wf_fragments_'//trim(intc(ik))//'_'//trim(intc(ib_grp))
         call cp_directory(trim(core_io_path_save)//"/SAVE/"//trim(fragment_name),trim(more_io_path)//"/SAVE",ierr)
         fragment_name='ns.wf_fragments_'//trim(intc(ik))//'_'//trim(intc(ib_grp))
         call cp_directory(trim(core_io_path_save)//"/SAVE/"//trim(fragment_name),trim(more_io_path)//"/SAVE",ierr)
       enddo
     enddo
     !
     call msg('l','done')
     !
   endif
   !
   ! Rotating wavefunctions by components
   !
   allocate(wf_disk(2,wf_nb_io,wf_ncx_new,n_spin))
   call mem_est("wf_disk",(/size(wf_disk)/),(/SP/))
   !
   if(wf_ncx_save/=wf_ncx_new) then
     !
     call live_timing('WFs comp. reduction',k_save%nibz)
     !
     call io_control(ACTION=OP_WR_CL,COM=NONE,SEC=(/1/),ID=ID)
     io_err=ioWF(ID)
     !
     ACTION_=OP_APP_WR_CL
     if(Fragmented_IO) ACTION_=OP_WR_CL
     !
     do ik=1,k_save%nibz
       !
       ik_save=kpoints_map(1,ik)
       !
       do ib_grp=1,wf_nb_io_groups
         !
         wf_disk=0._SP
         !
         nb1=wf_nb_io*(ib_grp-1)+1
         nb2=wf_nb_io*ib_grp
         !
         nkibz=k_save%nibz
         wf_ncx=wf_ncx_save
         core_io_path=core_io_path_save
         call WF_load(wf_ng,1,(/nb1,nb2/),(/ik_save,ik_save/),space='C',title='',impose_free_and_alloc=.true.)
         !
         do ib=nb1,nb2
           do i_spin=1,n_spin
             i_wf=wf_state(ib,ik_save,i_spin)
             wf_disk(1,ib,:wf_nc_k_new(ik),i_spin)=  real(wf(:wf_nc_k_save(ik_save),i_wf))
             wf_disk(2,ib,:wf_nc_k_new(ik),i_spin)= aimag(wf(:wf_nc_k_save(ik_save),i_wf))
           enddo
         enddo
         !
         nkibz=k%nibz
         wf_ncx=wf_ncx_new
         core_io_path=more_io_path
         call io_control(ACTION=ACTION_,COM=NONE,SEC=(/ik+1,1/),ID=ID)
         io_err=ioWF(ID,wf=wf_disk)
         !
       enddo
       !
       call live_timing(steps=1)
     enddo
     !
     call live_timing()
     !
   endif
   !
   n_steps=(k%nibz-k_save%nibz)
   if(n_steps>0) call live_timing('WFs rotation',n_steps)
   !
   allocate(ic_rot_table(wf_ncx_new))
   allocate(wf_tmp(wf_ncx_new,n_spin))
   !
   ACTION_=OP_APP_WR_CL
   if(Fragmented_IO) ACTION_=OP_WR_CL
   !
   do ik=k_save%nibz+1,k%nibz 
     !      
     ik_save=kpoints_map(1,ik)
     is=kpoints_map(2,ik)
     !
     ic_rot_table=-1
     do ic=1,wf_nc_k_save(ik_save)
       ig_rot=g_rot(is,wf_igk_save(ic,ik_save))
       do ic_rot=1,wf_nc_k_new(ik)
         ig=wf_igk_new(ic_rot,ik)
         if(ig==ig_rot) exit
       enddo
       ic_rot_table(ic)=ic_rot
     enddo
     !
     do ib_grp=1,wf_nb_io_groups
       !
       wf_disk=0._SP
       !
       nb1=wf_nb_io*(ib_grp-1)+1
       nb2=wf_nb_io*ib_grp
       !
       nkibz=k_save%nibz
       wf_ncx=wf_ncx_save
       core_io_path=core_io_path_save
       call WF_load(wf_ng,1,(/nb1,nb2/),(/ik_save,ik_save/),space='C',title='',impose_free_and_alloc=.true.)
       !
       do ib=nb1,nb2
         !
         wf_tmp=(0._SP,0._SP)
         !
         do i_spin=1,n_spin
           !
           i_wf=wf_state(ib,ik_save,i_spin)
           !
           forall(ic=1:wf_nc_k_save(ik_save)) wf_tmp(ic_rot_table(ic),i_spin)=wf(ic,i_wf)
           !
         enddo
         !
         forall(ic=1:wf_nc_k_new(ik)) wf_tmp(ic,:)=matmul(spin_sop(:,:,is),wf_tmp(ic,:))
         if ( S_contains_TR(is) ) wf_tmp=conjg(wf_tmp)
         !
         wf_disk(1,ib,1:wf_ncx_new,:)= real(wf_tmp(:wf_ncx_new,:))
         wf_disk(2,ib,1:wf_ncx_new,:)= aimag(wf_tmp(:wf_ncx_new,:))
         !
       enddo
       !
       nkibz=k%nibz
       wf_ncx=wf_ncx_new
       core_io_path=more_io_path
       call io_control(ACTION=ACTION_,COM=NONE,SEC=(/ik+1,1/),ID=ID)
       io_err=ioWF(ID,wf=wf_disk)
       !
     enddo
     !
     if(n_steps>0) call live_timing(steps=1)
     !
   enddo
   !
   if(n_steps>0) call live_timing()
   !
   deallocate(ic_rot_table)
   deallocate(wf_tmp)
   !
 case(2)
   !
   call msg('s',':: WFs reduction...')
   !
   call io_control(ACTION=OP_WR_CL,COM=NONE,SEC=(/1/),ID=ID)
   io_err=ioWF(ID)
   !
   allocate(wf_disk(2,wf_nb_io,wf_ncx_new,n_spin))
   call mem_est("wf_disk",(/size(wf_disk)/),(/SP/))
   !
   ACTION_=OP_APP_WR_CL
   if(Fragmented_IO) ACTION_=OP_WR_CL
   !
   do ik=1,k%nibz
     !
     ik_save=kpoints_map(1,ik)
     !
     do ib_grp=1,wf_nb_io_groups
       !
       wf_disk=0._SP
       !
       nb1=wf_nb_io*(ib_grp-1)+1
       nb2=wf_nb_io*ib_grp
       !
       nkibz=k_save%nibz
       wf_ncx=wf_ncx_save
       core_io_path=core_io_path_save
       call WF_load(wf_ng,1,(/nb1,nb2/),(/ik_save,ik_save/),space='C',title='',impose_free_and_alloc=.true.)
       !
       do ib=nb1,nb2
         do i_spin=1,n_spin
           i_wf=wf_state(ib,ik_save,i_spin)
           wf_disk(1,ib,1:wf_ncx_new,i_spin)= real(wf(1:wf_ncx_new,i_wf))
           wf_disk(2,ib,1:wf_ncx_new,i_spin)= aimag(wf(1:wf_ncx_new,i_wf))
         enddo
       enddo
       !
       nkibz=k%nibz
       wf_ncx=wf_ncx_new
       core_io_path=more_io_path
       call io_control(ACTION=ACTION_,COM=NONE,SEC=(/ik+1,1/),ID=ID)
       io_err=ioWF(ID,wf=wf_disk)
       !
     enddo
     !
   enddo
   !
   call msg('l','done')
   !
 end select
 ! 
 if(Fragmented_IO) then
   call io_control(ACTION=OP_WR_CL,COM=NONE,SEC=(/1/),ID=ID)
   io_err=ioWF(ID)
 endif
 !
 core_io_path=core_io_path_save
 !
 !
 ! CLEAN
 !=======
 call E_reset(E_new)
 call WF_free()
 deallocate(wf_disk)
 call mem_est("wf_disk")
 !
end subroutine
!
subroutine map_energies(E,E_new,k,kpoints_map)
 !
 use electrons,           ONLY:levels,E_reset,n_sp_pol
 use R_lattice,           ONLY:bz_samp
 !
 implicit none
 !
 type(levels), intent(in)     :: E
 type(levels), intent(out)    :: E_new
 type(bz_samp), intent(in)    :: k
 integer,       intent(in)    :: kpoints_map(2,k%nibz)
 !
 ! Work Space
 !
 integer :: ik
 !
 call E_reset(E_new)
 E_new%nb=E%nb
 E_new%nk=k%nibz
 allocate(E_new%E(E%nb,k%nibz,n_sp_pol))
 !
 do ik=1,k%nibz
   E_new%E(:,ik,:)=E%E(:,kpoints_map(1,ik),:)
 enddo
 !
end subroutine map_energies
